from sklearn.naive_bayes import GaussianNB
from sklearn.model_selection import train_test_split
import data_set
import prettytable
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# 转换成轴，3*n矩阵向量转置为n*3
def to_axios(data):
    x = []
    y = []
    z = []
    for e in data:
        x.append(e[0])
        y.append(e[1])
        z.append(e[2])
    return x, y, z


# 画分布图
def get_fig(data, target):
    train_group0 = []
    train_group1 = []
    train_group2 = []

    # 按照标签分组
    for i in range(0, len(target)):
        if data[i][0] > 1:
            continue
        if target[i] == "低":
            train_group0.append(data[i])
        elif target[i] == "中":
            train_group1.append(data[i])
        else:
            train_group2.append(data[i])

    fig = plt.figure()
    ax = Axes3D(fig)
    x0, y0, z0 = to_axios(train_group0)
    x1, y1, z1 = to_axios(train_group1)
    x2, y2, z2 = to_axios(train_group2)
    ax.scatter(x0, y0, z0, c="green", label="低")
    ax.scatter(x1, y1, z1, c="yellow", label="中")
    ax.scatter(x2, y2, z2, c="red", label="高")

    ax.set_zlabel('Z', fontdict={'color': 'red'})
    ax.set_ylabel('Y', fontdict={'color': 'red'})
    ax.set_xlabel('X', fontdict={'color': 'red'})
    plt.show()
    return


# 画图表
def get_chart(data, target, result):
    # 表头
    chart = prettytable.PrettyTable(["编号", "dim1", "dim2", "dim3", "实际结果", "预测结果"])
    for i in range(len(data)):
        if data[i][0] > 1:
            continue
        row = [i]
        row += data[i]
        row.append(target[i])
        row.append(result[i])
        chart.add_row(row)
    print(chart)
    return str(chart)



set = data_set.get_data_set()
x_train, x_test, y_train, y_test = train_test_split(set["data"], set["target"], test_size=0.3, random_state=1)
clf = GaussianNB()
clf.fit(x_train, y_train)
print(clf.score(x_train, y_train))
print(clf.score(x_test, y_test))
get_fig(data=x_train, target=y_train)
get_fig(data=x_test, target=y_test)
get_fig(data=x_test, target=list(clf.predict(x_test)))
get_fig(data=x_train, target=list(clf.predict(x_train)))
chart = get_chart(x_train, y_train, list(clf.predict(x_train)))
# 把表格写到文件中
# 开新表
with open("chart.txt", "w", encoding="utf-8") as f:
    f.write(chart)
f.close()
chart = get_chart(x_test, y_test, list(clf.predict(x_test)))
with open("chart.txt", "a", encoding="utf-8") as f:
    f.write(chart)
f.close()

